import os
import pandas as pd
from sklearn.utils import shuffle

label_list = []
image_list = []

image_dir = '/home/aistudio/ColorImage'
labels_dir = '/home/aistudio/Gray_Label'
csv_dir = '/home/aistudio/work/data_list'
expose_file_list = os.path.join(csv_dir, 'expose_file.txt')

"""
   ColorImage/
     ColorImage_road02/
       ColorImage/
         Record002/
           Camera 5/
            ...
           Camera 6
         Record003
       ....
     ColorImage_road03
     ColorImage_road04
   Gray_Label/
     Label_road02/
      Label
       Record002/
         Camera 5/
          ...
         Camera 6
       Record003
       ....
     Label_road03
     Label_road04     

"""
# ColorImage
image_list = []
label_list = []
expose_file = open(expose_file_list, 'r')
expose_images = [line.strip() for line in expose_file.readlines()]
expose_file.close()
print(expose_images)
for s1 in os.listdir(image_dir):
    if not s1.startswith('.'):
        image_sub_dir1 = os.path.join(image_dir, s1, 'ColorImage')
        label_sub_dir1 = os.path.join(labels_dir, s1.replace('ColorImage', 'Label'), 'Label')
        # print(image_sub_dir1)
        # print(label_sub_dir1)
        for s2 in os.listdir(image_sub_dir1):
            if not s2.startswith('.'):
                image_sub_dir2 = os.path.join(image_sub_dir1, s2)
                # print('image_sub_dir2', image_sub_dir2)
                label_sub_dir2 = os.path.join(label_sub_dir1, s2)
                # print('label_sub_dir2', label_sub_dir2)
                # camera
                for s3 in os.listdir(image_sub_dir2):
                    if not s3.startswith('.'):
                        # new_s3 = s3.replace(' ', '_')
                        # os.rename(os.path.join(image_sub_dir2, s3), os.path.join(image_sub_dir2, new_s3))
                        image_sub_dir3 = os.path.join(image_sub_dir2, s3)
                        # print(image_sub_dir3)
                        # os.rename(os.path.join(label_sub_dir2, s3), os.path.join(label_sub_dir2, new_s3))
                        label_sub_dir3 = os.path.join(label_sub_dir2, s3)
                        # print(label_sub_dir3)
                    for s4 in os.listdir(image_sub_dir3):
                        if not s4.startswith('.'):
                            s44 = s4.replace('.jpg', '_bin.png')
                            image_sub_dir4 = os.path.join(image_sub_dir3, s4)
                            # print(image_sub_dir4)
                            if image_sub_dir4 not in expose_images:
                                image_list.append(image_sub_dir4)
                                label_sub_dir4 = os.path.join(label_sub_dir3, s44)
                                # print(label_sub_dir4)
                                label_list.append(label_sub_dir4)
total_length = len(image_list)
print('total_length', total_length)
eighth_part = int(0.8 * total_length)
nine_part = int(0.9 * total_length)
all_pd = pd.DataFrame({'image': image_list, 'label': label_list})
all_shuffle = shuffle(all_pd)

train_dataset = all_shuffle[:eighth_part]
val_dataset = all_shuffle[eighth_part:nine_part]
test_dataset = all_shuffle[nine_part:]

train_dataset.to_csv(os.path.join(csv_dir, 'train.csv'), sep=' ', header=True, index=False)
val_dataset.to_csv(os.path.join(csv_dir, 'val.csv'), sep=' ', header=True, index=False)
test_dataset.to_csv(os.path.join(csv_dir, 'test.csv'), sep=' ', header=True, index=False)
